{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0c504ffb",
   "metadata": {},
   "source": [
    "# **Instructions**:\n",
    "\n",
    "### 1 - Open your game that you want to perform detections\n",
    "### 2 - In the game window, get the name of it's title bar \n",
    "### 3 - Update the variable \"window_name\" with the game title bar name\n",
    "### 4 - Run all cells to start detecting objects using your trained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91800ec7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import win32gui, win32ui, win32con\n",
    "from PIL import Image\n",
    "from time import sleep\n",
    "import cv2 as cv\n",
    "import os\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f9c9c67",
   "metadata": {},
   "outputs": [],
   "source": [
    "class WindowCapture:\n",
    "    w = 0\n",
    "    h = 0\n",
    "    hwnd = None\n",
    "\n",
    "    def __init__(self, window_name):\n",
    "        self.hwnd = win32gui.FindWindow(None, window_name)\n",
    "        if not self.hwnd:\n",
    "            raise Exception('Window not found: {}'.format(window_name))\n",
    "\n",
    "        window_rect = win32gui.GetWindowRect(self.hwnd)\n",
    "        self.w = window_rect[2] - window_rect[0]\n",
    "        self.h = window_rect[3] - window_rect[1]\n",
    "\n",
    "        border_pixels = 8\n",
    "        titlebar_pixels = 30\n",
    "        self.w = self.w - (border_pixels * 2)\n",
    "        self.h = self.h - titlebar_pixels - border_pixels\n",
    "        self.cropped_x = border_pixels\n",
    "        self.cropped_y = titlebar_pixels\n",
    "\n",
    "    def get_screenshot(self):\n",
    "        wDC = win32gui.GetWindowDC(self.hwnd)\n",
    "        dcObj = win32ui.CreateDCFromHandle(wDC)\n",
    "        cDC = dcObj.CreateCompatibleDC()\n",
    "        dataBitMap = win32ui.CreateBitmap()\n",
    "        dataBitMap.CreateCompatibleBitmap(dcObj, self.w, self.h)\n",
    "        cDC.SelectObject(dataBitMap)\n",
    "        cDC.BitBlt((0, 0), (self.w, self.h), dcObj, (self.cropped_x, self.cropped_y), win32con.SRCCOPY)\n",
    "\n",
    "        signedIntsArray = dataBitMap.GetBitmapBits(True)\n",
    "        img = np.fromstring(signedIntsArray, dtype='uint8')\n",
    "        img.shape = (self.h, self.w, 4)\n",
    "\n",
    "        dcObj.DeleteDC()\n",
    "        cDC.DeleteDC()\n",
    "        win32gui.ReleaseDC(self.hwnd, wDC)\n",
    "        win32gui.DeleteObject(dataBitMap.GetHandle())\n",
    "\n",
    "        img = img[...,:3]\n",
    "        img = np.ascontiguousarray(img) \n",
    "            \n",
    "        return img\n",
    "\n",
    "    def generate_image_dataset(self):\n",
    "        if not os.path.exists(\"images\"):\n",
    "            os.mkdir(\"images\")\n",
    "        while(True):\n",
    "            img = self.get_screenshot()\n",
    "            im = Image.fromarray(img[..., [2, 1, 0]])\n",
    "            im.save(f\"./images/img_{len(os.listdir('images'))}.jpeg\")\n",
    "            sleep(1)\n",
    "    \n",
    "    def get_window_size(self):\n",
    "        return (self.w, self.h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f22385d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImageProcessor:\n",
    "    W = 0\n",
    "    H = 0\n",
    "    net = None\n",
    "    ln = None\n",
    "    classes = {}\n",
    "    colors = []\n",
    "\n",
    "    def __init__(self, img_size, cfg_file, weights_file):\n",
    "        np.random.seed(42)\n",
    "        self.net = cv.dnn.readNetFromDarknet(cfg_file, weights_file)\n",
    "        self.net.setPreferableBackend(cv.dnn.DNN_BACKEND_OPENCV)\n",
    "        self.ln = self.net.getLayerNames()\n",
    "        self.ln = [self.ln[i-1] for i in self.net.getUnconnectedOutLayers()]\n",
    "        self.W = img_size[0]\n",
    "        self.H = img_size[1]\n",
    "        \n",
    "        with open('yolov4-tiny/obj.names', 'r') as file:\n",
    "            lines = file.readlines()\n",
    "        for i, line in enumerate(lines):\n",
    "            self.classes[i] = line.strip()\n",
    "        \n",
    "        # If you plan to utilize more than six classes, please include additional colors in this list.\n",
    "        self.colors = [\n",
    "            (0, 0, 255), \n",
    "            (0, 255, 0), \n",
    "            (255, 0, 0), \n",
    "            (255, 255, 0), \n",
    "            (255, 0, 255), \n",
    "            (0, 255, 255)\n",
    "        ]\n",
    "        \n",
    "\n",
    "    def proccess_image(self, img):\n",
    "\n",
    "        blob = cv.dnn.blobFromImage(img, 1/255.0, (416, 416), swapRB=True, crop=False)\n",
    "        self.net.setInput(blob)\n",
    "        outputs = self.net.forward(self.ln)\n",
    "        outputs = np.vstack(outputs)\n",
    "        \n",
    "        coordinates = self.get_coordinates(outputs, 0.5)\n",
    "\n",
    "        self.draw_identified_objects(img, coordinates)\n",
    "\n",
    "        return coordinates\n",
    "\n",
    "    def get_coordinates(self, outputs, conf):\n",
    "\n",
    "        boxes = []\n",
    "        confidences = []\n",
    "        classIDs = []\n",
    "\n",
    "        for output in outputs:\n",
    "            scores = output[5:]\n",
    "            \n",
    "            classID = np.argmax(scores)\n",
    "            confidence = scores[classID]\n",
    "            if confidence > conf:\n",
    "                x, y, w, h = output[:4] * np.array([self.W, self.H, self.W, self.H])\n",
    "                p0 = int(x - w//2), int(y - h//2)\n",
    "                boxes.append([*p0, int(w), int(h)])\n",
    "                confidences.append(float(confidence))\n",
    "                classIDs.append(classID)\n",
    "\n",
    "        indices = cv.dnn.NMSBoxes(boxes, confidences, conf, conf-0.1)\n",
    "\n",
    "        if len(indices) == 0:\n",
    "            return []\n",
    "\n",
    "        coordinates = []\n",
    "        for i in indices.flatten():\n",
    "            (x, y) = (boxes[i][0], boxes[i][1])\n",
    "            (w, h) = (boxes[i][2], boxes[i][3])\n",
    "\n",
    "            coordinates.append({'x': x, 'y': y, 'w': w, 'h': h, 'class': classIDs[i], 'class_name': self.classes[classIDs[i]]})\n",
    "        return coordinates\n",
    "\n",
    "    def draw_identified_objects(self, img, coordinates):\n",
    "        for coordinate in coordinates:\n",
    "            x = coordinate['x']\n",
    "            y = coordinate['y']\n",
    "            w = coordinate['w']\n",
    "            h = coordinate['h']\n",
    "            classID = coordinate['class']\n",
    "            \n",
    "            color = self.colors[classID]\n",
    "            \n",
    "            cv.rectangle(img, (x, y), (x + w, y + h), color, 2)\n",
    "            cv.putText(img, self.classes[classID], (x, y - 10), cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)\n",
    "        cv.imshow('window',  img)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94a699cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell to initiate detections using the trained model.\n",
    "\n",
    "window_name = \"YOUR_GAME_WINDOW_NAME_HERE\"\n",
    "cfg_file_name = \"./yolov4-tiny/yolov4-tiny-custom.cfg\"\n",
    "weights_file_name = \"yolov4-tiny-custom_last.weights\"\n",
    "\n",
    "wincap = WindowCapture(window_name)\n",
    "improc = ImageProcessor(wincap.get_window_size(), cfg_file_name, weights_file_name)\n",
    "\n",
    "while(True):\n",
    "    \n",
    "    ss = wincap.get_screenshot()\n",
    "    \n",
    "    if cv.waitKey(1) == ord('q'):\n",
    "        cv.destroyAllWindows()\n",
    "        break\n",
    "\n",
    "    coordinates = improc.proccess_image(ss)\n",
    "    \n",
    "    for coordinate in coordinates:\n",
    "        print(coordinate)\n",
    "    print()\n",
    "    \n",
    "    # If you have limited computer resources, consider adding a sleep delay between detections.\n",
    "    # sleep(0.2)\n",
    "\n",
    "print('Finished.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03ef282a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
